// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <string>

#include "pollux/buffer/string_view_buffer_holder.h"
#include "pollux/common/memory/memory.h"
#include "pollux/common/strings/string_view.h"

namespace kumo::pollux {
    namespace {
        std::string nonInlinedString() {
            return std::string(13, 'a');
        }

        std::string inlinedString() {
            return "a";
        }
    } // namespace

    class StringViewBufferHolderTest : public testing::Test {
    protected:
        static void SetUpTestCase() {
            memory::MemoryManager::initialize({});
        }

        StringViewBufferHolder makeHolder() {
            return StringViewBufferHolder(pool_.get());
        }

        std::shared_ptr<memory::MemoryPool> pool_{
                memory::memoryManager()->addLeafPool()};
    };

    TEST_F(StringViewBufferHolderTest, inlinedStringViewDoesNotCopyToBuffer) {
        auto holder = makeHolder();

        std::string value = inlinedString();
        StringView originalStringViewValue(value);
        ASSERT_TRUE(originalStringViewValue.isInline());

        auto ownedStringView = holder.getOwnedValue(originalStringViewValue);
        ASSERT_TRUE(ownedStringView.isInline());

        ASSERT_EQ(originalStringViewValue, ownedStringView);

        auto buffers = holder.moveBuffers();
        ASSERT_EQ(0, buffers.size());
    }

    TEST_F(StringViewBufferHolderTest, nonInlinedStringViewCopiesToBuffer) {
        auto holder = makeHolder();

        std::string value = nonInlinedString();
        StringView originalStringViewValue(value);
        ASSERT_FALSE(originalStringViewValue.isInline());

        auto ownedStringView = holder.getOwnedValue(originalStringViewValue);
        ASSERT_FALSE(ownedStringView.isInline());

        ASSERT_EQ(originalStringViewValue, ownedStringView);

        auto buffers = holder.moveBuffers();
        ASSERT_EQ(1, buffers.size());

        ASSERT_EQ(ownedStringView.data(), buffers.at(0)->as<char>());
    }

    TEST_F(StringViewBufferHolderTest, moreBuffersAreUsedWhenCurrentBufferFills) {
        // We're setting this high enough such that a new buffer allocation should be
        // needed by the StringViewBufferHolder. The min allocation is set to 8KB, but
        // the memory allocator allocates more than the requested size, so just
        // setting to much higher to guarantee the first buffer allocated is filled.
        size_t approxBytesToAdd = 1024 * 20;

        size_t nonInlinedBytesAdded = 0;
        auto holder = makeHolder();
        std::vector<std::string> addedStrings;

        std::vector<StringView> ownedStringViews;

        auto inlineIndex = [](size_t index) { return index % 2 == 0; };

        while (nonInlinedBytesAdded < approxBytesToAdd) {
            if (inlineIndex(addedStrings.size())) {
                addedStrings.push_back(inlinedString());
            } else {
                // Just to make sure the tests don't pass if only the memory location at
                // offset 0 was being used, change the suffix of the string each time
                addedStrings.push_back(
                        nonInlinedString() + std::to_string(addedStrings.size()));
                nonInlinedBytesAdded += addedStrings.back().size();
            }

            ownedStringViews.push_back(
                    holder.getOwnedValue(StringView(addedStrings.back())));
        }

        auto buffers = holder.moveBuffers();

        // Double checking that approxBytesToAdd was large enough to cause more than
        // one buffer to be allocated.
        ASSERT_GT(buffers.size(), 1);

        // Intentionally doing the comparison at the end in case memory was getting
        // overridden
        for (size_t i = 0; i < ownedStringViews.size(); ++i) {
            ASSERT_EQ(inlineIndex(i), ownedStringViews.at(i).isInline());
            ASSERT_EQ(addedStrings.at(i), ownedStringViews.at(i).str());
            ASSERT_NE(addedStrings.at(i).data(), ownedStringViews.at(i).data());
        }
    }

    TEST_F(StringViewBufferHolderTest, stateIsClearedAfterStringsAreMoved) {
        auto holder = makeHolder();
        std::string value = nonInlinedString();

        holder.getOwnedValue(StringView(value));

        auto firstMoved = holder.moveBuffers();
        ASSERT_EQ(1, firstMoved.size());

        auto secondMoved = holder.moveBuffers();
        ASSERT_EQ(0, secondMoved.size());
    }

    TEST_F(StringViewBufferHolderTest, getOwnedValueCanBeCalledWithIntegerType) {
        auto holder = makeHolder();
        ASSERT_EQ(42, holder.getOwnedValue(42));
        ASSERT_EQ(0, holder.moveBuffers().size());
    }

    TEST_F(StringViewBufferHolderTest, getOwnedValueCanBeCalledWithStringType) {
        const char *buf = "abcdefghijklmnopqrstuvxz";
        StringView result;

        auto holder = makeHolder();
        ASSERT_EQ(0, holder.buffers().size());

        {
            std::string str = buf;
            result = holder.getOwnedValue(str);
        }
        // Make sure `str` is already destructed at this point.
        ASSERT_EQ(StringView(buf), result);
        ASSERT_EQ(1, holder.buffers().size());
    }

    TEST_F(
            StringViewBufferHolderTest,
            getOwnedValueCanBeCalledWithStringPieceType) {
        const char *buf = "abcdefghijklmnopqrstuvxz";
        StringView result;
        melon::StringPiece piece;

        auto holder = makeHolder();
        ASSERT_EQ(0, holder.buffers().size());

        {
            std::string str = buf;
            piece = str;
            result = holder.getOwnedValue(piece);
        }

        // `str` is already destructed and piece is invalid.
        ASSERT_EQ(StringView(buf), result);
        ASSERT_EQ(1, holder.buffers().size());
    }

    TEST_F(StringViewBufferHolderTest, buffersCopy) {
        auto holder = makeHolder();
        holder.getOwnedValue(nonInlinedString());

        auto buffers = holder.buffers();
        EXPECT_EQ(1, buffers.size());

        auto firstMoved = holder.moveBuffers();
        EXPECT_EQ(1, firstMoved.size());
        EXPECT_EQ(1, buffers.size());

        auto secondMoved = holder.moveBuffers();
        EXPECT_EQ(0, secondMoved.size());
        EXPECT_EQ(1, firstMoved.size());
        EXPECT_EQ(1, buffers.size());

        auto secondBuffers = holder.buffers();
        EXPECT_EQ(0, secondBuffers.size());
        EXPECT_EQ(0, secondMoved.size());
        EXPECT_EQ(1, firstMoved.size());
        EXPECT_EQ(1, buffers.size());

        EXPECT_EQ(buffers.back()->as<uint8_t>(), firstMoved.back()->as<uint8_t>());
        EXPECT_NE(&buffers, &firstMoved);
    }

} // namespace kumo::pollux
